{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "id": "QM4ftkqa_rFO",
    "outputId": "28d7f50c-2344-42f3-ced1-338fd44b263f"
   },
   "outputs": [],
   "source": [
    "# Install the latest version of ParlAI from source.\n",
    "# This is important, otherwise using Blenderbot 2.0 will not work.\n",
    "!pip uninstall -q parlai -y \n",
    "!pip install -q  --progress-bar off git+https://github.com/facebookresearch/ParlAI.git \n",
    "\n",
    "# Install fairseq. This is required for Blenderbot 2\n",
    "!pip install -q fairseq\n",
    "\n",
    "# Clone the server and install its requirements\n",
    "!git clone https://github.com/JulesGM/ParlAI_SearchEngine.git\n",
    "!pip install -q  --progress-bar off -r ParlAI_SearchEngine/requirements.txt"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "sNiDF1HBAD5I"
   },
   "outputs": [],
   "source": [
    "import os\n",
    "import multiprocess\n",
    "import shlex\n",
    "import subprocess\n",
    "import threading\n",
    "import time\n",
    "\n",
    "# Using port 8080 doesn't work on Colab\n",
    "HOST = \"0.0.0.0:1111\"\n",
    "# Change the following as neede\n",
    "PATH_TO_SEARCH_SERVER = \"./ParlAI_SearchEngine/search_server.py\"\n",
    "\n",
    "assert os.path.exists(PATH_TO_SEARCH_SERVER), (\n",
    "    f\"Incorrect path {PATH_TO_SEARCH_SERVER}\"\n",
    ")\n",
    "\n",
    "command = [\"python\", \"-u\", shlex.quote(PATH_TO_SEARCH_SERVER), \n",
    "           \"serve\", \"--host\", HOST]\n",
    "command_str = \" \".join(command)\n",
    "p = subprocess.Popen(\n",
    "    command, \n",
    "    stderr=subprocess.STDOUT,\n",
    "    stdout=subprocess.PIPE,\n",
    ")\n",
    "\n",
    "# Wait a bit before the next cell to let a lot of the potential errors happen.\n",
    "time.sleep(1)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "E02KURghGWiW"
   },
   "outputs": [],
   "source": [
    "# Test If the server crashed. \n",
    "# You can rerun this cell to check if the server crashed. \n",
    "if p.poll() is not None: \n",
    "    print(p.communicate()[0].decode())\n",
    "    # If this says that the adress is already used, then it is very likely that\n",
    "    # you already started a server with a different Process object than p."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "id": "EkTzYpOOfoTa",
    "outputId": "bc8411c9-4d70-4482-eb98-3172459cde9f"
   },
   "outputs": [],
   "source": [
    "# Make sure we get a reasonable response.\n",
    "# Should be a list of json objects with a title, a content and an url field.\n",
    "!curl -X POST $HOST -d \"q=baseball&n=1\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "id": "D12RVzIFAF5C",
    "outputId": "e414497d-a5b9-41cc-9577-1b91f472d73c"
   },
   "outputs": [],
   "source": [
    "!python -m parlai interactive --model-file zoo:blenderbot2/blenderbot2_400M/model --search_server $HOST"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "f2IuHcvcD_sd"
   },
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "colab": {
   "collapsed_sections": [],
   "name": "Blenderbot2SearchEngineExample.py",
   "provenance": []
  },
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.8.5"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 1
}
